import logging

import orjson
from fastapi import Request
from fastapi.responses import JSONResponse, ORJSONResponse
from starlette.middleware.base import BaseHTTPMiddleware, RequestResponseEndpoint

from app.core.context import get_request_id
from app.schemas.response import Response as CustomResponse
from app.schemas.response import ResponseCode

logger = logging.getLogger(__name__)


class ResponseMiddleware(BaseHTTPMiddleware):
    async def dispatch(self, request: Request, call_next: RequestResponseEndpoint):
        try:
            response = await call_next(request)

            if isinstance(response, (JSONResponse, ORJSONResponse)):
                if response.status_code >= 400:
                    return ORJSONResponse(
                        content=CustomResponse.error().model_dump(),
                        status_code=200,
                    )

                try:
                    if isinstance(response.body, bytes):
                        data = orjson.loads(response.body)
                    else:
                        data = response.body
                except:
                    data = response.body.decode() if isinstance(response.body, bytes) else response.body

                return ORJSONResponse(content=CustomResponse.success(data=data).model_dump(), status_code=200)

            return response
        except Exception as e:
            logger.error(f"Response middleware error - ID: {get_request_id()} - Error: {str(e)}")
            return ORJSONResponse(
                content=CustomResponse.error(message="Internal Server Error").model_dump(), status_code=200
            )
